#include "disk.h"

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/utsname.h>
#include <sys/ioctl.h>

#include <linux/fs.h>

#include <scsi/sg.h>
#include <scsi/scsi_ioctl.h>

#include <unistd.h>
#include <string.h>
#include <stdarg.h>
#include <fcntl.h>

#include <functional>

namespace {
  // remove first and last flag
  std::string trim(const std::string &str, const std::string &flag) {
    int b = str.find_first_not_of(flag);
    int e = str.find_last_not_of(flag);
    return std::move(str.substr(b, e - b + 1));
  }

  // Return true if STR starts with PREFIX
  inline bool str_starts_with(const char * str, const char * prefix) {
    return !strncmp(str, prefix, strlen(prefix));
  }

  // return (v)sprintf() formatted std::string
  std::string vstrprintf(const char * fmt, va_list ap) {
    char buf[512];
    vsnprintf(buf, sizeof(buf), fmt, ap);
    buf[sizeof(buf)-1] = 0;
    return buf;
  }

  // return (v)sprintf() formatted std::string
  std::string strprintf(const char * fmt, ...) {
    va_list ap; va_start(ap, fmt);
    std::string str = vstrprintf(fmt, ap);
    va_end(ap);
    return str;
  }

  // callback be called every line. 
  // callback return true: continue, false: stop 
  bool file_each_line(const char *file, std::function<bool(const std::string &line)> callback, int max_line_len = 256) {
    if (!file || !*file || !callback || max_line_len <= 0)
      return false;
    FILE *f = ::fopen(file, "r");
    if (!f) return false;

    char *buf = (char*)malloc(max_line_len);
    if (!buf) {
      ::fclose (f);
      return false;
    }
    memset(buf, 0, max_line_len);

    while (fgets(buf, max_line_len, f)) {
      if (!callback(buf))
        break;
    }

    free(buf);
    ::fclose (f);
  }

  // read file all content
  bool file_content(const char *file, std::string &content) {
    if (!file || !*file)
      return false;
    FILE *f = ::fopen(file, "r");
    if (!f) return false;

    ::fseek(f, 0, SEEK_END);
    long epos = ::ftell(f);
    ::fseek(f, 0, SEEK_SET);
    long bpos = ::ftell(f);

    long len = epos - bpos;
    char *buf = (char*)malloc(len);
    if (!buf) {
      ::fclose (f);
      return false;
    }
    memset(buf, 0, len);

    ::fread(buf, sizeof(char*), len, f);
    content = buf;

    free(buf);
    ::fclose (f);
    return true;
  }

  // parser key and value
  // key=value,  key = value , key.value and so on. defalut is '='
  bool parse_k_v(const char *str, std::string &key, std::string &value, const char *flag = "=") {
    if (!str || !*str)
      return false;

    std::string src = str;
    auto pos = src.find(flag);
    if (pos == std::string::npos)
      return false;
    key = src.substr(0, pos);
    pos += strlen(flag);
    value = src.substr(pos, src.length() - pos);

    key = trim(key, " ");
    key = trim(key, "\n");
    value = trim(value, " ");
    value = trim(value, "\n");

    return true;
  }

  bool parse_file_v(const char *file, std::string &value) {
    if (!file || !*file)
      return false;
    if (!file_content(file, value))
      return false;

    value = trim(value, " ");
    value = trim(value, "\n");

    return true;
  }

  // Return kernel release as integer ("2.6.31" -> 206031)
  unsigned int get_kernel_release() {
    struct utsname u;
    if (uname(&u))
      return 0;
    unsigned x = 0, y = 0, z = 0;
    if (!(sscanf(u.release, "%u.%u.%u", &x, &y, &z) == 3
          && x < 100 && y < 100 && z < 1000             ))
      return 0;
    return x * 100000 + y * 1000 + z;
  }

  // get size in bytes
  bool disk_size(int fd, uint64_t *bytes) {
    auto ver = get_kernel_release();
    if (ver >= 206000
      || (ver >= 204018 && ver < 205000)) {
      return ioctl(fd, BLKGETSIZE64, bytes) >= 0;
    }

    uint32_t size;
    if (ioctl(fd, BLKGETSIZE, &size) >= 0) {
      *bytes = size;
      return true;
    }

    return false;
  }

}

namespace {
  /// Error (number,message) pair
  static device::sErrorInfo err_;
  
  /// Set last error number and message.
  /// Printf()-like formatting is supported.
  bool set_err(int no, const char * msg, ...);

  /// Set last error info struct.
  bool set_err(const device::sErrorInfo & err);

  /// Clear last error info.
  void clear_err();

  const char * type_to_string(device::Disk::eType type);

  // Read USB ID from /sys file
  static bool read_id(const std::string & path, unsigned short & id);
  // Get USB bridge ID for "sdX" or "sgN"
  static bool get_usb_id(const char * name, unsigned short & vendor_id,
    unsigned short & product_id, unsigned short & version);

  // Check for SCSI host proc_name "hpsa"
  static bool is_hpsa(const char * name);

  // err
  std::shared_ptr<device::Disk> missing_option(const char * opt);

  // Return disk object for disk 'name'.
  // Return 0 on error.
  std::shared_ptr<device::Disk> autodetect_disk(const char * name);
}

namespace device {
  bool Disk::scan_disk(std::set<std::string> &disks) {
    disks.clear();

    FILE *f = ::fopen("/proc/partitions", "r");
    if (!f) return false;

    char line[100], dev[100];
    int max, min, size;
    char *s;
    while (fgets(line, sizeof(line), f)) {
      if (sscanf(line, " %u %u %u %[^\n]"
        , &max, &min, &size, dev) != 4) {
        continue;
      }
      
      if (!disk_exist(dev))
        continue;

      std::string tmp = "/dev/";
      tmp += dev;
      disks.insert(std::move(tmp));
    }

    ::fclose(f);
    return true;
  }

  // Return disk object for disk 'name'.
  // Return 0 on error.
  std::shared_ptr<Disk> Disk::get_disk(const char * name) {
    clear_err();
    if (!name || !*name) {
      set_err(EINVAL, "param name is null or empty.");
      return nullptr;
    }
    return autodetect_disk(name);
  }

  // Get last error info struct.
  const sErrorInfo& Disk::get_err() { return err_; }

  bool Disk::disk_exist(const char * name) {
    const char *test_name = name;
    // Remove the leading /dev/... if it's there
    static const char dev_prefix[] = "/dev/";
    if (str_starts_with(test_name, dev_prefix))
      test_name += strlen(dev_prefix);

    char buf[128] = { 0 };
    snprintf(buf, 127
      , "/sys/block/%s/uevent"
      , test_name);

    bool bminor = false;
    bool bname = false;
    bool btype = false;
    std::string k,v;
    file_each_line(buf, [&bminor, &bname, &btype, &k, &v, &test_name](const std::string &line) -> bool {
      if (!parse_k_v(line.c_str(), k, v, "="))
        return true;
      if (k == "MINOR" && v == "0")
        bminor = true;
      if (k == "DEVNAME" && v == test_name)
        bname = true;
      if (k == "DEVTYPE" && v == "disk")
        btype = true;
      return true;
    });
    return (bminor && bname && btype);
  }
}

namespace device {
  // Disk implement
  class DiskImpl : public Disk {
    public:
      DiskImpl(const char * dev_name, eType type, int flags, int retry_flags = -1)
        : type_(type)
        , dev_(dev_name)
        , fd_(-1)
        , flags_(flags)
        , retry_flags_(retry_flags) {}
      virtual ~DiskImpl() { if (fd_ >= 0) ::close(fd_); }
    public: // for override
      // Return disk type
      const eType& get_type() const override { return type_; }

      // Return disk type string
      const char * get_type_str() const override { return type_to_string(type_); }

      // Return Device (path)name
      const char * get_dev() const override { return dev_.c_str(); }
    public:
      bool is_open() const override { return (fd_ >= 0); }
      bool open() override {
        if (is_open())
          return true;

        fd_ = ::open(dev_.c_str(), flags_);

        if (fd_ < 0 && errno == EROFS && retry_flags_ != -1)
          // Retry
          fd_ = ::open(dev_.c_str(), retry_flags_);

        if (fd_ < 0) {
          if (errno == EBUSY && (flags_ & O_EXCL))
            // device is locked
            return set_err(EBUSY,
              "The requested controller is used exclusively by another process!\n"
              "(e.g. smartctl or smartd)\n"
              "Please quit the impeding process or try again later...");
          return set_err((errno==ENOENT || errno==ENOTDIR) ? ENODEV : errno, strerror(errno));
        }

        if (fd_ >= 0) {
          // sets FD_CLOEXEC on the opened device file descriptor.  The
          // descriptor is otherwise leaked to other applications (mail
          // sender) which may be considered a security risk and may result
          // in AVC messages on SELinux-enabled systems.
          if (-1 == fcntl(fd_, F_SETFD, FD_CLOEXEC)) {
            // TODO: Provide an error printing routine in class smart_interface
            std::string msg = strprintf("fcntl(set  FD_CLOEXEC) failed, errno=%d [%s]\n", errno, strerror(errno));
            set_err(errno, msg.c_str());
          }
        }

        return true;
      }
      bool close() override {
        // ???
        int fd = fd_; fd_ = -1;
        if (::close(fd) < 0)
          return set_err(errno, strerror(errno));
        return true;
      }
      bool get_base_info(sBaseInfo &info) override {
        const char *test_name = dev_.c_str();
        // Remove the leading /dev/... if it's there
        static const char dev_prefix[] = "/dev/";
        if (str_starts_with(test_name, dev_prefix))
          test_name += strlen(dev_prefix);

        char buf[128] = { 0 };

        snprintf(buf, 127
          , "/sys/block/%s/device/vendor"
          , test_name);
        parse_file_v(buf, info.vendor);

        snprintf(buf, 127
          , "/sys/block/%s/device/model"
          , test_name);
        parse_file_v(buf, info.product);

        std::string tmp;
        snprintf(buf, 127
          , "/sys/block/%s/size"
          , test_name);
        parse_file_v(buf, tmp);
        info.sectors = std::stoull(tmp);

        disk_size(fd_, &info.size);

        return true;
      }
      bool get_inquiry(sInquiry &info) override {
        std::string msg = strprintf("The get_inquiry function of %s have not been implemented yet.", get_type_str());
        return set_err(0, msg.c_str());
      }
    protected:
      eType  type_;
      std::string dev_;   //< Device (path)name
      int fd_;            //< filedesc, -1 if not open.
      int flags_;         //< Flags for ::open()
      int retry_flags_;   ///< Flags to retry ::open(), -1 if no retry
  };

  // ATA support
  class AtaDisk : public DiskImpl {
    public:
      explicit AtaDisk(const char * dev_name)
        : DiskImpl(dev_name, eType::ata, O_RDONLY | O_NONBLOCK) {}
      virtual ~AtaDisk() { close(); }
    public:
  };

  // Standard SCSI support
  class ScsiDisk : public DiskImpl {
    public:
      explicit ScsiDisk(const char * dev_name)
        : DiskImpl(dev_name, eType::scsi, O_RDONLY | O_NONBLOCK) {}
      virtual ~ScsiDisk() { close(); }
    public:
      bool get_inquiry(sInquiry &info) override {
        struct sg_io_hdr *hdr = init_io_hdr(32 * 256);
        if (!hdr) {
          set_err(1, "init_io_hdr error.");
          return false;
        }

        unsigned char cdb[6];
        /* set the cdb format */
        cdb[0] = 0x12; /*This is for Inquery*/
        cdb[1] = 0 & 0xff;
        cdb[2] = 0 & 0xff;
        cdb[3] = 0;
        cdb[4] = 0xff;
        cdb[5] = 0; /*For control filed, just use 0*/

        hdr->dxfer_direction = SG_DXFER_FROM_DEV;
        hdr->cmdp = cdb;
        hdr->cmd_len = 6;

        if (ioctl(fd_, SG_IO, hdr) < 0) {
          set_err(1, "call ioctl error.");
          destroy_io_hdr(hdr);
          return false;
        }

        read_data(hdr, info);

        destroy_io_hdr(hdr);
        return true;
      }
    protected:
      void read_data(struct sg_io_hdr *hdr, sInquiry &info) {
        char * buffer = (char*)hdr->dxferp;
        info.vendor.assign(&buffer[8], 8);
        info.product.assign(&buffer[16], 16);
        info.version.assign(&buffer[32], 4);
      }
    protected:
      struct sg_io_hdr * init_io_hdr(unsigned int xfer_len, unsigned int sense_len = 32) {
        if (xfer_len <= 0 || sense_len <= 0)
          return nullptr;

        struct sg_io_hdr * p_hdr = (struct sg_io_hdr *)malloc(sizeof(struct sg_io_hdr));
        memset(p_hdr, 0, sizeof(struct sg_io_hdr));
        if (!p_hdr)
          return nullptr;

        p_hdr->interface_id = 'S'; /* this is the only choice we have! */
        p_hdr->flags = SG_FLAG_LUN_INHIBIT; /* this would put the LUN to 2nd byte of cdb*/
        p_hdr->dxferp = malloc(xfer_len);
        p_hdr->dxfer_len = xfer_len;
        p_hdr->sbp = (unsigned char*)malloc(sense_len);
        p_hdr->mx_sb_len = sense_len;

        if (!p_hdr->dxferp || !p_hdr->sbp) {
          destroy_io_hdr(p_hdr);
          return nullptr;
        }

        memset(p_hdr->dxferp, 0, p_hdr->dxfer_len);
        memset(p_hdr->sbp, 0, p_hdr->mx_sb_len);

        return p_hdr;
      }
      void destroy_io_hdr(struct sg_io_hdr * p_hdr) {
        if (!p_hdr) return;

        if (p_hdr->dxferp)
          delete ((char*)p_hdr->dxferp);
        if (p_hdr->sbp)
          delete ((char*)p_hdr->sbp);
        free(p_hdr);
      }

      // 序列号
      // https://www.xuebuyuan.com/993888.html
  };

  // PMC AacRAID support
  // LSI MegaRAID support
  // CCISS RAID support
  // AMCC/3ware RAID support
  // Areca RAID support
  // Marvell support
  // Highpoint RAID support
  // NVMe support
  class NvmeDisk : public DiskImpl {
    public:
      explicit NvmeDisk(const char * dev_name, unsigned nsid)
        : DiskImpl(dev_name, eType::nvme, O_RDONLY | O_NONBLOCK)
        , nsid_(nsid) {}
      virtual ~NvmeDisk() { close(); }
    public:
    protected:
      unsigned nsid_;
  };
}

namespace {
  bool set_err(int no, const char * msg, ...) {
    err_.no = no;
    if (!msg) {
      err_.msg.clear();
      return false;
    }

    va_list ap;
    va_start(ap, msg);
    err_.msg = vstrprintf(msg, ap);
    va_end(ap);
    return false;
  }
  bool set_err(const device::sErrorInfo & err) {
    err_.reset(err.no, err.msg.c_str());
    return false;
  }
  void clear_err() {
    err_.clear();
  }

  const char * type_to_string(device::Disk::eType type) {
    switch (type) {
    case device::Disk::eType::ata : return "ata";
    case device::Disk::eType::scsi : return "scsi";
    case device::Disk::eType::nvme : return "nvme";
    }
    return "unknown";
  }

  bool read_id(const std::string & path, unsigned short & id) {
    FILE * f = fopen(path.c_str(), "r");
    if (!f)
      return false;
    int n = -1;
    bool ok = (fscanf(f, "%hx%n", &id, &n) == 1 && n == 4);
    fclose(f);
    return ok;
  }
  bool get_usb_id(const char * name, unsigned short & vendor_id,
    unsigned short & product_id, unsigned short & version) {
    // Only "sdX" or "sgN" supported
    if (!(name[0] == 's' && (name[1] == 'd' || name[1] == 'g') && !strchr(name, '/')))
      return false;

    // Start search at dir referenced by symlink
    // "/sys/block/sdX/device" or
    // "/sys/class/scsi_generic/sgN"
    // -> "/sys/devices/.../usb*/.../host*/target*/..."
    std::string dir = strprintf("/sys/%s/%s%s",
      (name[1] == 'd' ? "block" : "class/scsi_generic"), name,
      (name[1] == 'd' ? "/device" : ""));

    // Stop search at "/sys/devices"
    struct stat st;
    if (stat("/sys/devices", &st))
      return false;
    ino_t stop_ino = st.st_ino;

    // Search in parent directories until "idVendor" is found,
    // fail if "/sys/devices" reached or too many iterations
    int cnt = 0;
    do {
      dir += "/..";
      if (!(++cnt < 10 && !stat(dir.c_str(), &st) && st.st_ino != stop_ino))
        return false;
    } while (access((dir + "/idVendor").c_str(), 0));

    // Read IDs
    if (!(   read_id(dir + "/idVendor", vendor_id)
          && read_id(dir + "/idProduct", product_id)
          && read_id(dir + "/bcdDevice", version)   ))
      return false;

    return true;
  }

  bool is_hpsa(const char * name) {
    char path[128];
    snprintf(path, sizeof(path), "/sys/block/%s/device", name);
    char * syshostpath = realpath(path, (char *)0);
    if (!syshostpath)
      return false;

    char * syshost = strrchr(syshostpath, '/');
    if (!syshost) {
      free(syshostpath);
      return false;
    }

    char * hostsep = strchr(++syshost, ':');
    if (hostsep)
      *hostsep = 0;

    snprintf(path, sizeof(path), "/sys/class/scsi_host/host%s/proc_name", syshost);
    free(syshostpath);
    int fd = open(path, O_RDONLY);
    if (fd < 0)
      return false;

    char proc_name[32];
    ssize_t n = read(fd, proc_name, sizeof(proc_name) - 1);
    close(fd);
    if (n < 4)
      return false;

    proc_name[n] = 0;
    if (proc_name[n - 1] == '\n')
      proc_name[n - 1] = 0;

    if (strcmp(proc_name, "hpsa"))
      return false;

    return true;
  }

  std::shared_ptr<device::Disk> missing_option(const char * opt) {
    set_err(EINVAL, "requires option '%s'", opt);
    return nullptr;
  }
  
  // Guess device type (ata or scsi) based on device name (Linux
  // specific) SCSI device name in linux can be sd, sr, scd, st, nst,
  // osst, nosst and sg
  std::shared_ptr<device::Disk> autodetect_disk(const char * name) {
    const char * test_name = name;

    // Dereference symlinks
    struct stat st;
    std::string pathbuf;
    if (!lstat(name, &st) && S_ISLNK(st.st_mode)) {
      char * p = realpath(name, (char *)0);
      if (p) {
        pathbuf = p;
        free(p);
        test_name = pathbuf.c_str();
      }
    }

    // Remove the leading /dev/... if it's there
    static const char dev_prefix[] = "/dev/";
    if (str_starts_with(test_name, dev_prefix))
      test_name += strlen(dev_prefix);

    // form /dev/h* or h*
    if (str_starts_with(test_name, "h"))
      return std::make_shared<device::AtaDisk>(name);

    // form /dev/ide/* or ide/*
    if (str_starts_with(test_name, "ide/"))
      return std::make_shared<device::AtaDisk>(name);

    // form /dev/s* or s*
    if (str_starts_with(test_name, "s")) {

      // Try to detect possible USB->(S)ATA bridge
      unsigned short vendor_id = 0, product_id = 0, version = 0;
      if (get_usb_id(test_name, vendor_id, product_id, version)) {
        /*const char * usbtype = get_usb_dev_type_by_id(vendor_id, product_id, version);
        if (!usbtype)
          return 0;

        // Kernels before 2.6.29 do not support the sense data length
        // required for SAT ATA PASS-THROUGH(16)
        if (!strcmp(usbtype, "sat") && get_kernel_release() < 206029)
          usbtype = "sat,12";

        // Return SAT/USB device for this type
        // (Note: linux_scsi_device::autodetect_open() will not be called in this case)
        return get_scsi_passthrough_device(usbtype, new linux_scsi_device(this, name, ""));*/

        set_err(0, "USB not implemented");
        return nullptr;
      }

      // Fail if hpsa driver
      if (is_hpsa(test_name))
        return missing_option("-d cciss,N");

      // No USB bridge or hpsa driver found, assume regular SCSI device
      return std::make_shared<device::ScsiDisk>(name);
    }

    // form /dev/scsi/* or scsi/*
    if (str_starts_with(test_name, "scsi/"))
      return std::make_shared<device::ScsiDisk>(name);

    // form /dev/bsg/* or bsg/*
    if (str_starts_with(test_name, "bsg/"))
      return std::make_shared<device::ScsiDisk>(name);

    // form /dev/ns* or ns*
    if (str_starts_with(test_name, "ns"))
      return std::make_shared<device::ScsiDisk>(name);

    // form /dev/os* or os*
    if (str_starts_with(test_name, "os"))
      return std::make_shared<device::ScsiDisk>(name);

    // form /dev/nos* or nos*
    if (str_starts_with(test_name, "nos"))
      return std::make_shared<device::ScsiDisk>(name);

    // form /dev/nvme* or nvme*
    if (str_starts_with(test_name, "nvme"))
      return std::make_shared<device::NvmeDisk>(name, 0); // use default nsid 

    // form /dev/tw[ael]* or tw[ael]*
    if (str_starts_with(test_name, "tw") && strchr("ael", test_name[2]))
      return missing_option("-d 3ware,N");

    // form /dev/cciss/* or cciss/*
    if (str_starts_with(test_name, "cciss/"))
      return missing_option("-d cciss,N");

    // we failed to recognize any of the forms
    set_err(0, "failed to recognize any of the forms");

    return std::make_shared<device::ScsiDisk>(name);

    return nullptr;
  }
}